/*
  Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

  This file is part of Sunrise.

  Sunrise is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 3 of the License, or
  (at your option) any later version.

  Sunrise is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Declarations of some handy vector operations.

// $Id$

#ifndef __vecops__
#define __vecops__

#include "mcrx-types.h"

#include <vector>
#include "blitz/array.h"
#include <blitz/array/stencil-et.h>
#include <blitz/array/stencil-et-macros.h>
#include <stdexcept>

namespace blitz {
 BZ_DECLARE_DIFF(forward1sum) {
   return (*A)+A.shift(1,dim);
 }

BZ_ET_STENCIL_DIFF(forward1sum, 0, 1);
};

namespace mcrx {

  /** Returns the magnitude of a 3-vector. */
  T_float mag(const vec3d& v);

  /** Takes a vector v and rotates it according to the rotation
      defined by taking z->newz and x->newx. If inverse is true, v is
      rotated according to the inverse transformation, ie taking
      newz->z, and newx->x. */
  vec3d rotate(const vec3d& v, 
	       const vec3d& newz, const vec3d& newx,
	       bool inverse=false);


  /** Returns a 1D blitz Array referencing the data in the vector. It
      is the user's responsibility to ensure that the vector does not
      go out of scope before the array. */
  template<typename T>
  inline blitz::Array<T,1> reference_vector(const std::vector<T>& v) {
    return blitz::Array<T,1>(const_cast<T*>(&v[0]), v.size(), 
			     blitz::neverDeleteData);
  }


  /** Takes a vector v and rotates it according to the rotation defined
      by taking z->newz and the x-axis onto the orthogonal projection of
      the x-axis. If the newz is nearly aligned with the x-axis then
      this transformation is singular and in that case the x-axis will
      go onto the negative z-axis to preserve smoothness of the
      transformation.  You should use this function if the azimuthal
      direction in the xy plane doesn't matter, otherwise you need to
      explicitly specify the new x direction with the trinary rotate
      function. */
  vec3d rotate(const vec3d& v, const vec3d& newz);

  /** Interpolates a point on the line (or power-law) segment between
      the specified points. */
  inline T_float interpolate_point(T_float x0, T_float y0,
				   T_float x1, T_float y1,
				   T_float x,
				   const bool logarithmic)
  {
    assert(x1>x0);
    assert(x>=x0);
    assert(x<=x1);
    
    if (!logarithmic) {
      return (y1-y0)/(x1-x0)*(x-x0)+y0;
    }
    else {
      if ((y0==0)||(y1==0))
	// power law is not defined, assume it's zero
	return 0;
      
      const T_float alpha = log(y1/y0)/log(x1/x0);
      return y0*pow(x/x0, alpha);
    }
  }


  /** Returns an array of y values associated with x values in newx,
      interpolated onto the function defined by the points in arrays x
      and y. The interpolation is done piecewise linear unless
      logarithmic is true, in which case it's done piecewise
      power-law. */
  array_1 interpolate_onto(const array_1& x, const array_1& y,
			   const array_1& newx,
			   bool logarithmic);


  /** Calculates the integral of a line segment
      connecting the two specified points. */
  inline T_float integrate_segment_lin(T_float x0, T_float y0,
				T_float x1, T_float y1)
  {
    assert (x1>=x0);
    return 0.5*(y1+y0)*(x1-x0);
  };

  /** Calculates the integral of a power-law segment
      connecting the two specified points. */
  inline T_float integrate_segment_log(T_float x0, T_float y0,
				T_float x1, T_float y1)
  {
    assert (x1>=x0);
    if (x0==x1)
      return 0;
  
    if ((y0==0) || (y1==0))
      // in this case a power-law segment is strictly undefined so we
      // set the integral to zero.
      return 0;
  
    assert (y0>0);
    assert (y1>0);
  
    T_float xr= x1/x0;
    T_float logxr = log(xr);
    T_float alpha = log(y1/y0)/logxr;
  
    if (std::abs(alpha+1)>1e-4)
      // use analytic integral formula
      return y0*x0*( pow(xr,alpha+1) -1 )/(alpha+1);
    else
      // use expansion when alpha is close to -1
      return y0*x0*( logxr + 0.5*logxr*logxr*(alpha+1) );
  }

  BZ_DECLARE_FUNCTION4(integrate_segment_lin)
  BZ_DECLARE_FUNCTION4(integrate_segment_log)
}


namespace blitz{
  BZ_DECLARE_STENCIL_OPERATOR2(integrate_segment_log_stencil, y, x)
  return mcrx::integrate_segment_log(*x, *y, x.shift(1,0), y.shift(1,0));
  BZ_END_STENCIL_OPERATOR

  BZ_DECLARE_STENCIL_OPERATOR2(integrate_segment_log_stencil_dim2, y, x)
  return mcrx::integrate_segment_log(*x, *y, x.shift(1,1), y.shift(1,1));
  BZ_END_STENCIL_OPERATOR

  BZ_ET_STENCIL2(integrate_segment_log_stencil, double , bzCC(BZ_PROMOTE(typename BZ_BLITZ_SCOPE(asExpr)<T1>::T_expr::T_numtype, typename BZ_BLITZ_SCOPE(asExpr)<T2>::T_expr::T_numtype)), shape(0), shape(1))
    BZ_ET_STENCIL2(integrate_segment_log_stencil_dim2, double , bzCC(BZ_PROMOTE(typename BZ_BLITZ_SCOPE(asExpr)<T1>::T_expr::T_numtype, typename BZ_BLITZ_SCOPE(asExpr)<T2>::T_expr::T_numtype)), shape(0,0), shape(0,1))
}


namespace mcrx {
  
  T_float integrate_between(const array_1& x,
			    const array_1& y,
			    T_float xmin, T_float xmax,
			    bool logarithmic,
			    int left_idx = -1,
			    int right_idx = -1);

  /** Calculates the integral of a quantity defined by arrays x and
      y. If logarithmic is false, the integration is done using the
      simple trapezoidal rule, ie the function is assumed to be
      piecewise linear between the points. If logarithmic is true,
      then the function is assumed to be a piecewise power law between
      the points, which is more appropriate in many
      circumstances. Special care is being taken to correctly handle
      segments where the approximation approaches 1/x. */
  template<typename T1, typename T2>
  inline T_float integrate_quantity (const blitz::ETBase<T1>& y,
				     const blitz::ETBase<T2>& x,
				     const bool logarithmic=false)
{
  assert(blitz::asExpr<T1>::T_expr::rank_==1);
  assert(blitz::asExpr<T2>::T_expr::rank_==1);

  if(logarithmic)
    return sum(blitz::integrate_segment_log_stencil(y.wrap(),x.wrap()));
  else
    // care must be taken to make sure that we don't call the stencil
    // operator itself. If the wrap() is not used, we get forward1sum
    // of the first element instead of the stencil...
    return 0.5*sum(blitz::forward1sum(y.wrap(), 0)*
		   blitz::forward11(x.wrap(), 0));
}

  /** Forwards to the expression integrate_quantity. Needed to convert
      array arguments to ETBase. */
  inline T_float integrate_quantity(const array_1& y, const array_1& x, 
				    bool logarithmic=false)
  {
    return integrate_quantity(y.wrap(), x.wrap(), logarithmic);
  }
  
  /// Return type for the integrate_quantities_lin function.
  template<typename T> struct RT_integrate_quantities_lin {

    typedef typename blitz::BzBinaryExprResult<
      blitz::Multiply,
      double,
      typename blitz::BzReductionResult<
	blitz::ReduceSum,
	1,
	typename blitz::BzBinaryExprResult<
	  blitz::Multiply,
	  typename blitz::BzStencilResult<
	    blitz::forward1sum_et,
	    T>::T_result,
	  typename blitz::BzStencilResult<
	    blitz::forward11_et,
	    typename blitz::BzIndexmapResult<
	      array_1,
	      1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
	      >::T_result
	    >::T_result
	  >::T_result
	>::T_result
      >::T_result T_result;
  };

  /** Calculates the integral over the second dimension for each of
      the quantities in the 2D array y using linear interpolation. The
      integration variable MUST be the second index. Because the
      return type is an expression, log and lin integrations must be
      distinct functions.
  */
  template<typename T>
  inline
  typename RT_integrate_quantities_lin<T>::T_result
  integrate_quantities_lin(const blitz::ETBase<T>& y,
			  const array_1& x)
  {
    using namespace blitz;
    return 0.5*sum(blitz::forward1sum(y.wrap(), 1)*
		   blitz::forward11(x(tensor::j),1), 
		   tensor::j);
  }

  /// Return type for the integrate_quantities_log function.
  template<typename T> struct RT_integrate_quantities_log {

    typedef typename blitz::BzIndexmapResult<
      array_1,
      1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0>::T_result T_map1;

    typedef typename blitz::BzReductionResult<
      blitz::ReduceSum,
      1,
      typename blitz::BzBinaryStencilResult<
	blitz::integrate_segment_log_stencil_dim2_et2,
	T,
	T_map1,
	BZ_PROMOTE(typename blitz::asExpr<T>::T_expr::T_numtype,
		   typename blitz::asExpr<T_map1>::T_expr::T_numtype)
	>::T_result
      >::T_result T_result;
  };

  /** Calculates the integral over the second dimension for each of
      the quantities in the 2D array y using log interpolation. The
      integration variable MUST be the second index. Because the
      return type is an expression, log and lin integrations must be
      distinct functions.
  */
  template<typename T>
  inline
  typename RT_integrate_quantities_log<T>::T_result
  integrate_quantities_log(const blitz::ETBase<T>& y, 
			   const array_1& x)
  {
    using namespace blitz;
    return sum(blitz::integrate_segment_log_stencil_dim2(y.wrap(),
							 x(tensor::j)),
	       tensor::j);
  }

  /** Calculates a vector of midpoints between the specified points,
      with special assumptions for the endpoints. */
  std::vector<mcrx::T_float> 
  calculate_middle(const std::vector<mcrx::T_float>& v);


  /** Calculates the delta vector indicating the width of the bins in
      the vector of interval midpoints supplied. This is handy
      e.g. when integrating quantities. The endpoint intervals are
      assumed to be symmetric about the midpoint. */
  template<typename T>
  blitz::Array<T, 1> delta_quantity(const blitz::Array<T, 1>& v) {
    const int n = v.size();
    blitz::Array<T, 1> delta (n);
    if (n>1) {
      blitz::Range r (1, n - 2);
      // calculate delta using stencil
      delta (r) = .5*(v (r+ 1) -v (r- 1));
      //delta (r) = blitz::central12n(v, blitz::firstDim);
      // Edge bins have half the width
      delta (0) = 0.5*(v (1) -v (0));
      delta (n- 1) = 0.5*(v (n- 1) -v (n- 2));
    }
    else {
      // if there's only one bin, we default to bin width one (this
      // makes it easier to deal with situations where the "bin"
      // really is not a d-quantity but a quantity
      delta=1.0;
    }
    return delta;
  }

  /** Calculates the delta vector indicating the width of the bins in
      the vector of interval midpoints supplied. This is handy
      e.g. when integrating quantities. The endpoint intervals are
      assumed to be symmetric about the midpoint. */
  template<typename T>
  blitz::Array<T, 1> delta_quantity(const std::vector<T>& v) {
    const blitz::Array<T, 1> a (reference_vector(v));
    return delta_quantity(a);
  }

  /** Function to resample a vector from an old x onto a new array of
      x values. It tries to do this intelligently by averaging the
      data over the new bins assuming linear interpolation. Note that
      unlike the subsample function below, it does *not* preserve the
      integral. If downsampling to lower resolutions, narrow features
      will be lost. For that reason, using subsample is
      recommended. */
  mcrx::array_1 resample(const array_1& data, const array_1& old_x,
			 const array_1& new_x, bool allow_outside,
			 bool logarithmic);

  /** Function to subsample a vector from an old x onto a new
      (smaller) array of x values. It does this so that the integral
      of the vector is preserved in each bin. It also is invariant to
      resampling onto the same x-values repeatedly.  Subsampling to
      preserve the integral only works if the integral is interpreted
      as a series of rectangles. If we try to come up with a scheme
      that preserves the integral as a series of piecewise power-law
      integrals, we get the unstable global optimization problem we
      had in the python code. Thus, this function has no option for
      logarithmic operation. It's called subsample, because if you try
      to upsample with it, it will still preserve the integral but
      artifacts from the linear interpolation will be obvious. */
  mcrx::array_1 subsample(const array_1& data, const array_1& old_x,
			  const array_1& new_x);


  /** Create an array of N logarithmically spaced set of values
      between xmin and xmax. */
  mcrx::array_1 logspace(T_float xmin, T_float xmax, int N);

}; // end namespace mcrx

#endif
